import os
import asyncio
import threading
import shutil

from django.conf import settings
import django
if not settings.configured:
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "codevaluation.settings")
    django.setup()

import subprocess
import sys
import json
import os
import re
import aiofiles
from urllib.parse import parse_qs
from channels.generic.websocket import AsyncWebsocketConsumer
from django.utils import timezone



class TrainingConsumer(AsyncWebsocketConsumer):
    max_epochs = None

    async def connect(self):
        await self.accept()

    async def receive(self, text_data):
        try:
            data = json.loads(text_data)

            if data['type'] == 'file_upload':
                await self.handle_file(data)
            elif data['type'] == 'parameters':
                await self.handle_parameters(data)
            elif data['type'] == 'reupload_file':
                await self.handle_reupload(data)

        except Exception as e:
            print("处理错误:", str(e))
            await self.send(json.dumps({
                'type': 'error',
                'message': f'处理错误: {str(e)}'
            }))

    async def handle_file(self, data):
        save_dir ="./model/open/"
        save_dir = os.path.normpath(save_dir)  # 规范化路径（去除多余的..和\）
        file_path = os.path.join(save_dir, data['name'])
        if os.path.exists(save_dir):
            shutil.rmtree(save_dir)
        os.makedirs(save_dir, exist_ok=True)
        try:
            # 写入文件
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(data['content'])

            # 验证文件
            file_size = os.path.getsize(file_path)

            # 读取前几行内容
            with open(file_path, 'r', encoding='utf-8') as f:
                print("文件前3行:")
                for _ in range(3):
                    line = f.readline()
                    if not line: break
                    print(line.strip())

            await self.send(json.dumps({
                'type': 'status',
                'message': f'文件 {data["name"]} 上传成功'
            }))

        except Exception as e:
            error_msg = f'文件保存失败: {str(e)}'
            print(error_msg)
            await self.send(json.dumps({
                'type': 'error',
                'message': error_msg
            }))

    async def handle_parameters(self, data):
        """处理训练参数"""
        max_epochs = data.get('epochs')
        print("\n=== 训练参数 ===")
        print("训练轮数:", max_epochs)
        print("批次比例大小:", data.get('batch_size_ratio'))
        print("序列长度:", data.get('squence_len'))
        print("学习率:", data.get('learning_rate'))
        await self.send(json.dumps({
            'type': 'status',
            'message': '参数接收成功，正在训练中……'
        }))

    async def handle_reupload(self, data):
    # 计算相对路径：向上返回两级（codevaluation → BIGCHUANG），然后进入 BIGCHUANGFILE
        save_dir = "./model/open"
        save_dir = os.path.normpath(save_dir)  # 规范化路径（去除多余的..和\）
        filename = data.get('filename')

        if not filename:
            await self.send(json.dumps({
                'type': 'error',
                'message': '未提供要删除的文件名'
            }))
            return

        file_path = os.path.join(save_dir, filename)

        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                await self.send(json.dumps({
                    'type': 'status',
                    'message': f'文件 {filename} 已成功删除'
                }))
                print(f"已删除文件: {file_path}")
            else:
                await self.send(json.dumps({
                    'type': 'status',
                    'message': f'文件 {filename} 不存在，无需删除'
                }))
        except Exception as e:
            error_msg = f'删除文件失败: {str(e)}'
            print(error_msg)
            await self.send(json.dumps({
                'type': 'error',
                'message': error_msg
            }))





class TrainingResultConsumer(AsyncWebsocketConsumer):
    trainresult_connections = set()

    async def connect(self):
        await self.accept()
        self.epoch_count = 0
        self.keep_alive = True
        self.training_started = False
        self.selected_epochs = set()
        self.training_parameters = None
        self.max_epochs = None  # 默认值

        query_params = parse_qs(self.scope["query_string"].decode())
        self.page_type = query_params.get("page", ["unknown"])[0]
        if self.page_type == "trainresult":
            TrainingResultConsumer.trainresult_connections.add(self)

    async def disconnect(self, close_code):
        self.keep_alive = False
        if self.page_type == "trainresult":
            TrainingResultConsumer.trainresult_connections.discard(self)

    async def receive(self, text_data):
        data = json.loads(text_data)
        try:
            if data['type'] == 'parameters':
                print("\n=== 收到训练参数 ===")
                print(f"参数内容: {data}")

                if not self.training_started:
                    self.training_started = True
                    self.training_parameters = data
                    self.max_epochs = int(data.get('epochs', 60))

                    print("[系统] 开始创建训练任务")
                    threading.Thread(target=self.run_train_script).start()



        except Exception as e:
            import traceback
            print(f"错误详情:\n{traceback.format_exc()}")

    def run_train_script(self):
        import os
        env = os.environ.copy()
        env["PYTHONUNBUFFERED"] = "1"

        print(f"[系统] 启动 train.py 脚本，训练轮数：{self.max_epochs}")
        print(self.training_parameters.get("epochs"))
        params_json = json.dumps(self.training_parameters)

        script_path = r'./model/FineSE.py'
        process = subprocess.Popen(
            [sys.executable,"-u", script_path, params_json],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=env
        )

        for line in iter(process.stdout.readline, ''):
            line = line.strip()
            if not line:
                continue
            print("[STDOUT]", line)
            try:

                if not line.startswith("METRICS_JSON:"):
                    continue  # 跳过非指标数据

                json_str = line[len("METRICS_JSON:"):].strip()
                metrics = json.loads(json_str)
                epoch = metrics.get('epoch', 0)
                mae = metrics.get('mae', 0)
                mmre = metrics.get('mmre', 0)
                mdae = metrics.get('mdae', 0)
                pred = metrics.get('pred', 0)

                print(f"调试：Epoch {epoch} 收到的指标数据：{metrics}")

                if epoch % 5 == 0:
                    asyncio.run(self.channel_layer.send(
                        self.channel_name,
                        {
                            "type": "send_to_frontend",
                            "text": json.dumps({
                                'type': 'epoch_ready',
                                'epoch': epoch
                            })
                        }
                    ))

                asyncio.run(self.channel_layer.send(
                    self.channel_name,
                    {
                        "type": "send_to_frontend",
                        "text": json.dumps({
                            'type': 'metrics_update',
                            'epoch': epoch,
                            'max_epochs': self.max_epochs,
                            'mae': mae,
                            'mmre': mmre,
                            'mdae': mdae,
                            'pred': pred,
                        })
                    }
                ))

            except Exception as e:
                print(f"[train.py 输出解析失败] {e}")
        # 最后发送训练完成
        asyncio.run(self.channel_layer.send(
            self.channel_name,
            {
                "type": "send_to_frontend",
                "text": json.dumps({
                    "type": "training_complete",
                    "max_epochs": self.max_epochs
                })
            }
        ))

    async def send_to_frontend(self, event):
        for conn in TrainingResultConsumer.trainresult_connections:
            await conn.send(event['text'])






# ⚠️ 临时变量，用于在两个 Consumer 之间共享上传的文件名（仅开发测试使用）
latest_uploaded_filename = None


class PredictConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        await self.accept()

    async def receive(self, text_data):
        global latest_uploaded_filename
        try:
            data = json.loads(text_data)

            if data['type'] == 'file_upload':
                await self.handle_file(data)
                # ✅ 记录上传的文件名（无路径）
                latest_uploaded_filename = data['name']
                print(f"[{timezone.now()}] 上传文件名记录为: {latest_uploaded_filename}")

            elif data['type'] == 'reupload_file':
                await self.handle_reupload(data)





        except Exception as e:
            print("处理错误:", str(e))
            await self.send(json.dumps({
                'type': 'error',
                'message': f'处理错误: {str(e)}'
            }))

    async def handle_file(self, data):
        save_dir ="./model/open/"
        os.makedirs(save_dir, exist_ok=True)

        file_path = os.path.join(save_dir, data['name'])
        if os.path.exists(save_dir):
            shutil.rmtree(save_dir)
        os.makedirs(save_dir, exist_ok=True)

        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(data['content'])

            print("文件前3行:")
            with open(file_path, 'r', encoding='utf-8') as f:
                for _ in range(3):
                    line = f.readline()
                    if not line: break
                    print(line.strip())

            await self.send(json.dumps({
                'type': 'status',
                'message': f'文件 {data["name"]} 上传成功'
            }))

        except Exception as e:
            await self.send(json.dumps({
                'type': 'error',
                'message': f'文件保存失败: {str(e)}'
            }))

    async def handle_reupload(self, data):
        save_dir ="./model/open"
        filename = data.get('filename')

        if not filename:
            await self.send(json.dumps({
                'type': 'error',
                'message': '未提供要删除的文件名'
            }))
            return

        file_path = os.path.join(save_dir, filename)

        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                await self.send(json.dumps({
                    'type': 'status',
                    'message': f'文件 {filename} 已成功删除'
                }))
            else:
                await self.send(json.dumps({
                    'type': 'status',
                    'message': f'文件 {filename} 不存在，无需删除'
                }))
        except Exception as e:
            await self.send(json.dumps({
                'type': 'error',
                'message': f'删除文件失败: {str(e)}'
            }))


class ResultConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        query_params = parse_qs(self.scope["query_string"].decode())
        self.page_type = query_params.get("page", ["unknown"])[0]

        if self.page_type == "result":
            await self.channel_layer.group_add("result_page", self.channel_name)

        await self.accept()

    async def receive(self, text_data):
        global latest_uploaded_filename
        try:
            data = json.loads(text_data)

            if data['type'] == 'start_prediction':
                username = data['username']
                model_id = data['model_id']
                print(f"开始预测 - 用户名: {username}, 模型ID: {model_id}")
                print("[DEBUG] 即将调用 predict.py 脚本")

                # 构建源路径和目标路径
                source_path = os.path.abspath(f'../BIGCHUANGFILE/MODELS/{model_id}.pt')
                target_dir = os.path.abspath('./model/model')
                target_path = os.path.join(target_dir, f'{model_id}.pt')

                #清空target_dir
                if os.path.exists(target_dir):
                    shutil.rmtree(target_dir)
                os.makedirs(target_dir, exist_ok=True)

                # 尝试复制模型文件
                try:
                    shutil.copyfile(source_path, target_path)
                    print(f"[DEBUG] 成功复制模型文件到 {target_path}")
                except Exception as e:
                    print(f"[ERROR] 无法复制模型文件: {e}")
                    # 你也可以选择这里 raise e 或 return 提前中止

                # 在子线程中运行 subprocess，实时读取输出
                def run_predict_script():
                    process = subprocess.Popen(
                        ['python', './model/predict.py'],
                        stdout=subprocess.PIPE,
                        stderr=subprocess.STDOUT,
                        text=True,
                        bufsize=1
                    )
                    for line in process.stdout:
                        print(f"[predict.py] {line.strip()}")
                    process.wait()

                await asyncio.to_thread(run_predict_script)

                print("[DEBUG] 脚本执行完成，准备处理结果")
                results, file_name, total_workload = await self.process_prediction_file(latest_uploaded_filename)

                await self.channel_layer.group_send(
                    "result_page",
                    {
                        "type": "send_analysis_results",
                        "results": results,
                        "file_name": file_name,
                        "total_workload": total_workload,
                    }
                )

        except Exception as e:
            import traceback
            error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
            await self.send_error("处理请求时发生错误，请查看服务器日志")
            print(f"[ERROR] 处理请求时发生错误:\n{error_msg}")


    async def send_analysis_results(self, event):
        await self.send(text_data=json.dumps({
            "type": "analysis_results",
            "file_name": event["file_name"],
            "results": event["results"],
            "total_workload": event["total_workload"]
        }))

    async def send_error(self, message):
        await self.send(text_data=json.dumps({
            "type": "error",
            "message": message
        }))

    @staticmethod
    async def process_prediction_file(csv_filename):
        base_name = os.path.splitext(csv_filename)[0]
        txt_filename = f"{base_name}.txt"
        total_workload = 0.0
        file_path = os.path.join( './model/pred_output/', txt_filename)

        async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
            content = await f.read()

        results = []
        entries = content.split('\n\n')
        line_number = 1

        for entry in entries:
            lines = entry.strip().split('\n')
            if len(lines) >= 2:
                task = lines[0].strip()
                workload_match = re.match(
                    r'^Predicted Workload:\s*\[([\d.]+)\]$',
                    lines[1].strip(),
                    re.IGNORECASE
                )
                if workload_match:
                    workload_value = float(workload_match.group(1))
                    total_workload += workload_value
                    results.append({
                        "id": str(line_number),
                        "title": task,
                        "content": f"Predicted Workload: {workload_match.group(1)}"
                    })
                    line_number += 1

        return results, txt_filename,round(total_workload, 3)
